#
# openlookup/engine/db.py
#
# Copyright 2008-2009 Helsinki Institute for Information Technology
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import with_statement
from contextlib import contextmanager

import time
import pickle
import logging
import sqlite3

from openlookup.engine.item import recreate_item
from openlookup.engine.identity import decode_identity
from openlookup.engine.signature import decode_signature

INFO_TABLE = """
create table if not exists info(
    k primary key,
    v)"""

DATA_TABLE = """
create table if not exists data(
    num integer primary key autoincrement,
    expiration integer,
    key unique,
    value,
    secret,
    priority default 0,
    type,
    pk,
    sig)"""

DATA_BY_EXPIRATION = """
create index if not exists by_expiration on data(expiration)"""

FILENAME = 'openlookup.db'

_undefined = object()

class DB(object):
    initialization = INFO_TABLE, DATA_TABLE, DATA_BY_EXPIRATION

    def __init__(self, filename=FILENAME, create=True, timeout=30, fast=False):
        self.db = sqlite3.Connection(filename, timeout=timeout,
                                     isolation_level=None)
        self.db.execute("pragma cache_size = 20000")
        if fast:
            self.db.execute("pragma synchronous = off")

        if create:
            self.db.execute("pragma auto_vacuum = 1")
            for statement in self.initialization:
                self.db.execute(statement)

            # XXX Temporary migration hack:
            if len(list(self.db.execute("pragma table_info(data)"))) == 6:
                logging.info('Adding type, pk, and sig columns to data table.')
                self.db.execute("alter table data add column type")
                self.db.execute("alter table data add column pk")
                self.db.execute("alter table data add column sig")

        self.transaction_in_progress = False

    def close(self):
        self.db.close()

    def get_info(self, key, default=_undefined):
        """Get configuration information."""
        cursor = self.db.execute("select v from info where k = ?", (key,))
        for value, in cursor:
            return pickle.loads(value)
        else:
            if default != _undefined:
                return default
            else:
                raise LookupError

    def set_info(self, key, value):
        """Set configuration information."""
        value = buffer(pickle.dumps(value))
        self.db.execute("replace into info values (?, ?)", (key, value))

    def clear_info(self, key):
        """Remove configuration information."""
        self.db.execute("delete from info where k = ?", (key,))

    def get(self, key):
        """
        Read an item from the database.

        The returned item may be a tombstone (value is None).  If no
        item at all is found, None is returned.
        """
        cursor = self.db.execute("select * from data"
                                 " where key = ?", (key,))
        for row in list(cursor):
            return item_from_sqlite(*row)
        else:
            return None

    def get_by_num(self, num):
        """
        Fetch an item from the database by number.
        """
        cursor = self.db.execute("select * from data"
                                 " where num = ?", (num,))
        for row in list(cursor):
            return item_from_sqlite(*row)
        else:
            return None

    def update(self, item, priority=1):
        """
        Add a newly created item to the database.

        If the item is not signed, the expiration time is adjusted if
        necessary to supersede a preexisting item.
        """
        if item.sig is not None:
            return self.put(item, priority)

        with self.transaction():
            old_item = self.get(item.key)
            if old_item is not None and item < old_item:
                item.expiration = old_item.expiration + 1
            self._put(item, priority)

    def put(self, item, priority=0):
        """
        Add an item to the database.
        """
        with self.transaction():
            old_item = self.get(item.key)
            if old_item is None or item > old_item:
                self._put(item, priority)

    def _put(self, item, priority):
        data = dict(expiration=item.expiration,
                    type=item.type,
                    key=item.key,
                    value=bytes_to_sqlite(item.value),
                    secret=bytes_to_sqlite(item.secret),
                    priority=priority)

        data['pk'] = buffer(item.pk.encode_public()) if item.pk else None
        data['sig'] = buffer(item.sig.encode()) if item.sig else None

        self.db.execute(
            "replace into"
            " data (expiration, key, value, secret, priority, type, pk, sig)"
            " values (:expiration, :key, :value, :secret, :priority,"
            "         :type, :pk, :sig)", data)

    def query(self, head, mark='', anchor=''):
        """
        Read non-tombstone items whose keys start with a specified prefix.

        This is a generator.

        The pertinent keys are of the form head + tail.  These are
        processed in circular lexicographic order by tail, running
        from mark (exclusive) to anchor (inclusive).
        """
        if mark < anchor:
            cursor = self.db.execute(
                "select * from data"
                " where key > ? and key <= ? and value is not null"
                " order by key", (head + mark, head + anchor))
            try:
                for row in cursor:
                    yield item_from_sqlite(*row)
            finally:
                cursor.close()
        else:
            cursor = self.db.execute(
                "select * from data"
                " where key > ? and value is not null"
                " order by key", (head + mark,))
            try:
                for row in cursor:
                    num, expiration, key = row[:3]
                    if not key.startswith(head):
                        break
                    yield item_from_sqlite(*row)
            finally:
                cursor.close()

            cursor = self.db.execute(
                "select * from data"
                " where key >= ? and key <= ? and value is not null"
                " order by key", (head, head + anchor))
            try:
                for row in cursor:
                    yield item_from_sqlite(*row)
            finally:
                cursor.close()

    def list_from(self, mark=0, priority=0):
        """
        Get a list of item numbers greater than _mark_.

        Only items with at least the specified priority are considered.
        """
        cursor = self.db.execute(
            "select num from data where num > ? and priority >= ?"
            " order by num asc", (mark, priority))
        return [num for num, in cursor]

    def list_to(self, mark, limit, priority=0):
        """
        Get a list of item numbers less than or equal to _mark_.

        At most _limit_ numbers are returned, in ascending order.  If
        more are present, the largest ones are returned.

        Only items with at least the specified priority are considered.
        """
        cursor = self.db.execute(
            "select num from data where num <= ? and priority >= ?"
            " order by num desc limit ?", (mark, priority, limit))
        nums = [num for num, in cursor]
        nums.reverse()
        return nums

    def get_mark(self):
        """
        Return the largest item number that has ever been present in the database.

        If the database is brand new, the return value is zero.  Item
        numbers are always greater than zero.
        """
        cursor = self.db.execute("select seq from sqlite_sequence"
                                 " where name='data'")
        for seq, in cursor:
            return seq
        else:
            return 0

    def clean(self, limit=1000):
        """
        Clean up expired entries.
        """
        with self.transaction():
            # Limiting how many entries are deleted on each call
            # avoids locking the database too long.  Unfortunately
            # sqlite3 by default does not enable using a limit clause
            # in a delete statement.
            now = int(time.time())
            cursor = self.db.execute(
                "select num from data"
                " where expiration < ? limit ?", (now, limit))
            nums = [num for num, in cursor]
            complete = (len(nums) < limit)
            if nums:
                logging.debug('Deleting %d entries.', len(nums))
                if complete:
                    self.db.execute("delete from data"
                                    " where expiration < ?", (now,))
                else:
                    for num in nums:
                        self.db.execute("delete from data"
                                        " where num = ?", (num,))
            return not complete

    @contextmanager
    def transaction(self):
        if self.transaction_in_progress:
            yield
        else:
            self.transaction_in_progress = True
            try:
                self.db.execute("begin immediate")
                try:
                    yield
                    self.db.execute("commit")
                except:
                    self.db.execute("rollback")
                    raise
            finally:
                self.transaction_in_progress = False

def item_from_sqlite(num, expiration, key, value, secret, priority,
                     type, pk, sig):
    key = str(key)              # Keys are hexadecimal.
    value = bytes_from_sqlite(value)
    secret = bytes_from_sqlite(secret)
    pk = bytes_from_sqlite(pk)
    sig = bytes_from_sqlite(sig)
    item = recreate_item(expiration=expiration, type=type,
                         key=key, value=value, secret=secret,
                         pk=pk, sig=sig)
    item.num = num
    item.priority = priority
    return item

def bytes_to_sqlite(x): return buffer(x) if isinstance(x, str) else x
def bytes_from_sqlite(x): return str(x) if isinstance(x, buffer) else x
